import "babel-polyfill";
import {forEach} from "../algorithm/iterator";

const [red, black] = [0, 1];

var Node = class {
    left = null;
    right = null;

    constructor({value, color, parent = null} = {}) {
        this.data = value;
        this.parent = parent;
        this.color = color;
    }
};

export default function(compare = (x, y) => x - y) {
    return class {
        _root = null;
        _size = 0;

        constructor() {
            this.add(...arguments);
        }

        get [Symbol.toStringTag]() {
            return "RBTree";
        }

        get size() {
            return this._size;
        }

        *[Symbol.iterator](direction = 1) {
            if (direction > 0) {
                var traverse = function* traverse(node) {
                    if (node) {
                        yield* traverse(node.left);
                        yield node.data;
                        yield* traverse(node.right);
                    }
                };
            } else {
                var traverse = function* traverse(node) {
                    if (node) {
                        yield* traverse(node.right);
                        yield node.data;
                        yield* traverse(node.left);
                    }
                };
            }
            yield* traverse(this._root);
        }

        _rotateLeft(node) {
            var right = node.right;
            var parent = node.parent;
            if (node.right = right.left) {
                right.left.parent = node;
            }
            right.left = node;
            right.parent = parent;
            if (parent) {
                if (node === parent.left) {
                    parent.left = right;
                } else {
                    parent.right = right;
                }
            } else {
                this._root = right;
            }
            node.parent = right;
        }

        _rotateRight(node) {
            var left = node.left;
            var parent = node.parent;
            if (node.left = left.right) {
                left.right.parent = node;
            }
            left.right = node;
            left.parent = parent;
            if (parent) {
                if (node === parent.right) {
                    parent.right = left;
                } else {
                    parent.left = left;
                }
            } else {
                this._root = left;
            }
            node.parent = left;
        }

        _insertFix(node) {
            var parent;
            while ((parent = node.parent) && parent.color === red) {
                let grandparent = parent.parent;
                if (parent === grandparent.left) {
                    let uncle = grandparent.right;
                    if (uncle && uncle.color === red) {
                        parent.color = black;
                        uncle.color = black;
                        grandparent.color = red;
                        node = grandparent;
                        continue;
                    }
                    if (node === parent.right) {
                        this._rotateLeft(parent);
                        [parent, node] = [node, parent];
                    }
                    parent.color = black;
                    grandparent.color = red;
                    this._rotateRight(grandparent);
                } else {
                    let uncle = grandparent.left;
                    if (uncle && uncle.color === red) {
                        parent.color = black;
                        uncle.color = black;
                        grandparent.color = red;
                        node = grandparent;
                        continue;
                    }
                    if (node === parent.left) {
                        this._rotateRight(parent);
                        [parent, node] = [node, parent];
                    }
                    parent.color = black;
                    grandparent.color = red;
                    this._rotateLeft(grandparent);
                }
            }
            this._root.color = black;
        }

        _eraseNode(node) {
            if (node.left && node.right) {
                let old = node;
                node = node.right;
                let left;
                while (left = node.left) {
                    node = left;
                }
                if (old.parent) {
                    if (old.parent.left === old) {
                        old.parent.left = node;
                    } else {
                        old.parent.right = node;
                    }
                } else {
                    this._root = node;
                }
                var child = node.right;
                var parent = node.parent;
                var node_color = node.color;
                if (parent === old) {
                    parent = node;
                } else {
                    if (child) {
                        child.parent = parent;
                    }
                    parent.left = child;
                    node.right = old.right;
                    old.right.parent = node;
                }
                node.parent = old.parent;
                node.color = old.color;
                node.left = old.left;
                old.left.parent = node;
            } else {
                var child = node.left || node.right;
                var parent = node.parent;
                var node_color = node.color;
                if (child) {
                    child.parent = parent;
                }
                if (parent) {
                    if (node === parent.left) {
                        parent.left = child;
                    } else {
                        parent.right = child;
                    }
                } else {
                    this._root = child;
                }
            }
            if (node_color === black) {
                this._eraseFix(child, parent);
            }
        }

        _eraseFix(node, parent) {
            while ((!node || node.color === black) && node !== this._root) {
                if (parent.left === node) {
                    let sibling = parent.right;
                    if (sibling.color === red) {
                        sibling.color = black;
                        parent.color = red;
                        this._rotateLeft(parent);
                        sibling = parent.right;
                    }
                    if ((!sibling.left || sibling.left.color === black) && (!sibling.right || sibling.right.color === black)) {
                        sibling.color = red;
                        node = parent;
                        parent = node.parent;
                    } else {
                        if (!sibling.right || sibling.right.color === black) {
                            sibling.left.color = black;
                            sibling.color = red;
                            this._rotateRight(sibling);
                            sibling = parent.right;
                        }
                        sibling.color = parent.color;
                        parent.color = black;
                        sibling.right.color = black;
                        this._rotateLeft(parent);
                        node = this._root;
                        break;
                    }
                } else {
                    let sibling = parent.left;
                    if (sibling.color === red) {
                        sibling.color = black;
                        parent.color = red;
                        this._rotateRight(parent);
                        sibling = parent.left;
                    }
                    if ((!sibling.right || sibling.right.color === black) && (!sibling.left || sibling.left.color === black)) {
                        sibling.color = red;
                        node = parent;
                        parent = node.parent;
                    } else {
                        if (!sibling.left || sibling.left.color === black) {
                            sibling.right.color = black;
                            sibling.color = red;
                            this._rotateLeft(sibling);
                            sibling = parent.left;
                        }
                        sibling.color = parent.color;
                        parent.color = black;
                        sibling.left.color = black;
                        this._rotateRight(parent);
                        node = this._root;
                        break;
                    }
                }
            }
            if (node) {
                node.color = black;
            }
        }

        find(value) {
            var node = this._root;
            while (node) {
                let comparison = compare(value, node.data);
                if (comparison === 0) {
                    return node;
                } else if (comparison < 0) {
                    node = node.left;
                } else {
                    node = node.right;
                }
            }
        }

        has(value) {
            return !!this.find(value);
        }

        add() {
            arguments::forEach(value => {
                if (this._root) {
                    let node = this._root;
                    let parent, dir;
                    while (node) {
                        let comparison = compare(value, node.data);
                        if (comparison === 0) {
                            return;
                        } else {
                            parent = node;
                            node = parent[dir = comparison < 0 ? "left" : "right"];
                        }
                    }
                    parent[dir] = node = new Node({value, color: red, parent});
                    this._insertFix(node);
                } else {
                    this._root = new Node({value, color: black});
                }
                ++this._size;
            });
            return this._size;
        }

        delete(value) {
            for (let node = this._root; node;) {
                let comparison = compare(value, node.data);
                if (comparison === 0) {
                    this._eraseNode(node);
                    --this._size;
                    return;
                } else if (comparison < 0) {
                    node = node.left;
                } else {
                    node = node.right;
                }
            }
            return value;
        }
    };
};
